import pdb
import os 
import numpy as np
from util import audio
import matplotlib.pyplot as plt
import pyworld as vocoder
from hparams import hparams as hp
import soundfile as sf

def visualize(lf0_file,mgc_file,bap_file,sample_id,dataset,res_dir):

    print("sample rate: {}".format(hp.sample_rate))
    norm_lf0_array = np.load(lf0_file)
    lf0_array = audio.f0_denormalize(norm_lf0_array)

    # fig = plt.figure()
    # plt.plot(lf0_array)
    # save_path = os.path.join(res_dir,dataset + '-' + sample_id + '-lf0.png')
    # plt.savefig(save_path)


    norm_mgc_array = np.load(mgc_file)
    mgc_array = audio.sp_denormalize(norm_mgc_array)
    pdb.set_trace()
    # print(norm_mgc_array.shape)

    # fig = plt.figure(figsize=(20, 5))
    # heatmap = plt.pcolor(norm_mgc_array.T)
    # fig.colorbar(mappable=heatmap)
    # plt.xlabel('Time(s)')
    # plt.ylabel('Mel Spectrum')
    # plt.tight_layout()
    # save_path = os.path.join(res_dir,dataset + '-' + sample_id + '-mgc.png')
    # plt.savefig(save_path)

    print("====================")
    # norm_bap_array = np.load(bap_file)
    # bap_array = audio.ap_denormalize(norm_bap_array,norm_lf0_array)
    original_bap_array = np.load(bap_file).astype(np.float64)
    
    fig = plt.figure(figsize=(20, 5))  
    heatmap = plt.pcolor(original_bap_array.T)
    np.savetxt(os.path.join(res_dir,"original_bap.txt"), original_bap_array.T,fmt='%f',delimiter=',')
    fig.colorbar(mappable=heatmap)
    plt.xlabel('Time(s)')
    plt.ylabel('AP')
    plt.tight_layout()
    save_path = os.path.join(res_dir,dataset + '-' + sample_id + '-original_bap.png')
    plt.savefig(save_path)
    print("original_bap: {}".format(original_bap_array.shape))


    # encoded_bap_array = vocoder.code_aperiodicity(original_bap_array,hp.sample_rate)
    # np.savetxt(os.path.join(res_dir,"encoded_bap.txt"), encoded_bap_array,fmt='%f',delimiter=',')

    # fig = plt.figure(figsize=(20, 5))  
    # heatmap = plt.pcolor(encoded_bap_array.T)
    # np.savetxt(os.path.join(res_dir,"encoded_bap.txt"), encoded_bap_array.T,fmt='%f',delimiter=',')
    # fig.colorbar(mappable=heatmap)
    # plt.xlabel('Time(s)')
    # plt.ylabel('AP')
    # plt.tight_layout()
    # save_path = os.path.join(res_dir,dataset + '-' + sample_id + '-encoded_bap.png')
    # plt.savefig(save_path)
    # print("encoded_bap: {}".format(encoded_bap_array.shape))
    

    # fig = plt.figure(figsize=(20, 5))  
    # # restored_bap_array = vocoder.decode_aperiodicity(encoded_bap_array,hp.sample_rate,fft_size = hp.fft_size)
    # restored_bap_array = vocoder.decode_aperiodicity(original_bap_array,hp.sample_rate,fft_size = hp.fft_size)
    restored_bap_array = audio.ap_denormalize(original_bap_array,norm_lf0_array)
    # heatmap = plt.pcolor(restored_bap_array.T)
    # np.savetxt(os.path.join(res_dir,"restored_bap.txt"), restored_bap_array.T,fmt='%f',delimiter=',')
    # fig.colorbar(mappable=heatmap)
    # plt.xlabel('Time(s)')
    # plt.ylabel('AP')
    # plt.tight_layout()
    # save_path = os.path.join(res_dir,dataset + '-' + sample_id + '-restored_bap.png')
    # plt.savefig(save_path)
    # print("restored_bap: {}".format(restored_bap_array.shape))

    wav = vocoder.synthesize(lf0_array, mgc_array, restored_bap_array, hp.sample_rate)
    path = os.path.join(res_dir,'synthesized.wav')
    # audio.save_wav(wav, path)
    sf.write(path, wav, 16000)
    # sf.write(path, wav, hp.sample_rate)
    pdb.set_trace()

if __name__ == "__main__":    
    dataset = 'debug_thchs30'
    sample_id = '00001'
    res_dir = 'visualize_res'
    
    if not os.path.exists(res_dir):
        os.makedirs(res_dir)
    lf0_file = os.path.join(dataset,'lf0-{}.npy'.format(sample_id))
    mgc_file = os.path.join(dataset,'mgc-{}.npy'.format(sample_id))
    bap_file = os.path.join(dataset,'bap-{}.npy'.format(sample_id))

    visualize(lf0_file,mgc_file,bap_file,sample_id,dataset,res_dir)
    